import torchvision
from torch import nn
from torch.utils.data import DataLoader

train_dataset = torchvision.datasets.CIFAR10(root='./CIFAR10',
                                             transform=torchvision.transforms.ToTensor(),
                                             train=True,
                                             download=True)

dataloader = DataLoader(dataset=train_dataset, batch_size=64)


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 3)

    def forward(self, x):
        return self.conv1(x)


mynet = MyNet()
for imgs, label in dataloader:
    print(imgs.shape, label)